# 导入必要的库
import numpy as np
import faiss
import pickle
from sklearn.neighbors import KNeighborsClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report
import os
import sys
from util import createXY, safe_normalize_L2, prepare_features

# 设置随机种子，保证结果可复现
np.random.seed(42)

# 内存优化参数
SAMPLE_RATIO = 0.2  # 仅使用20%的数据进行训练，可根据内存情况调整
RESIZE_SHAPE = (64, 64)  # 使用更小的图像尺寸
USE_FLOAT16 = True  # 使用float16节省内存

# 定义数据路径
DATA_PATH = 'cat_dog_data/cat_dog_data/data/train'

# 检查数据路径是否存在
print(f"正在检查数据路径: {os.path.abspath(DATA_PATH)}")
if not os.path.exists(DATA_PATH):
    print(f"错误: 数据路径不存在: {os.path.abspath(DATA_PATH)}")
    # 尝试寻找可能的替代路径
    print("尝试寻找可能的数据路径...")
    # 列出当前目录结构以帮助调试
    print("当前目录结构:")
    for root, dirs, files in os.walk('.', topdown=True):
        level = root.replace('.', '').count(os.sep)
        indent = ' ' * 4 * (level)
        print('{}{}/'.format(indent, os.path.basename(root)))
        subindent = ' ' * 4 * (level + 1)
        for f in files:
            print('{}{}'.format(subindent, f))
    sys.exit(1)

# 检查数据路径中的内容
print(f"数据路径内容: {os.listdir(DATA_PATH)}")

# 检查是否存在保存的数据文件
if os.path.exists('X.pkl') and os.path.exists('y.pkl'):
    print("正在加载已处理的数据...")
    # 加载保存的数据
    with open('X.pkl', 'rb') as f:
        X = pickle.load(f)
    with open('y.pkl', 'rb') as f:
        y = pickle.load(f)
    
    # 检查加载的数据是否为空
    if isinstance(X, list):
        print(f"加载的X是列表，长度: {len(X)}")
        X = np.array(X)
    else:
        print(f"加载的X形状: {X.shape}")
        
    if isinstance(y, list):
        print(f"加载的y是列表，长度: {len(y)}")
        y = np.array(y)
    else:
        print(f"加载的y形状: {y.shape}")
    
    # 如果加载的数据为空，重新创建数据
    if len(X) == 0:
        print("警告: 加载的数据为空，将重新创建数据...")
        os.remove('X.pkl')
        os.remove('y.pkl')
        X, y = createXY(train_folder=DATA_PATH, method='flat', resize_shape=RESIZE_SHAPE)
        with open('X.pkl', 'wb') as f:
            pickle.dump(X, f)
        with open('y.pkl', 'wb') as f:
            pickle.dump(y, f)
else:
    print("正在创建并处理数据...")
    # 创建数据 - 使用更小的图像尺寸
    X, y = createXY(train_folder=DATA_PATH, method='flat', resize_shape=RESIZE_SHAPE)
    
    # 保存处理后的数据，以便下次直接使用
    with open('X.pkl', 'wb') as f:
        pickle.dump(X, f)
    with open('y.pkl', 'wb') as f:
        pickle.dump(y, f)

# 确保 X 和 y 是 NumPy 数组
if isinstance(X, list):
    print("将 X 从列表转换为 NumPy 数组...")
    X = np.array(X)
if isinstance(y, list):
    print("将 y 从列表转换为 NumPy 数组...")
    y = np.array(y)

# 检查数据形状
print(f"数据特征形状: {X.shape}")
print(f"数据标签形状: {y.shape}")

# 确保数据不为空
if len(X) == 0:
    print("错误详情:")
    print(f"- 当前工作目录: {os.getcwd()}")
    print(f"- 数据路径: {os.path.abspath(DATA_PATH)}")
    print(f"- 数据路径是否存在: {os.path.exists(DATA_PATH)}")
    if os.path.exists(DATA_PATH):
        print(f"- 数据目录内容: {os.listdir(DATA_PATH)}")
        # 检查是否有子目录（cat和dog）
        has_cat_dir = os.path.exists(os.path.join(DATA_PATH, 'cat'))
        has_dog_dir = os.path.exists(os.path.join(DATA_PATH, 'dog'))
        print(f"- 包含cat目录: {has_cat_dir}")
        print(f"- 包含dog目录: {has_dog_dir}")
        # 检查文件格式
        if has_cat_dir:
            cat_files = os.listdir(os.path.join(DATA_PATH, 'cat'))
            print(f"- cat目录文件数量: {len(cat_files)}")
            if len(cat_files) > 0:
                print(f"- 示例cat文件名: {cat_files[:3]}")
        if has_dog_dir:
            dog_files = os.listdir(os.path.join(DATA_PATH, 'dog'))
            print(f"- dog目录文件数量: {len(dog_files)}")
            if len(dog_files) > 0:
                print(f"- 示例dog文件名: {dog_files[:3]}")
    raise ValueError("数据集为空，请检查数据加载过程")

# 采样部分数据以减少内存使用
if SAMPLE_RATIO < 1.0:
    print(f"采样 {int(SAMPLE_RATIO*100)}% 的数据以减少内存使用...")
    sample_size = int(len(X) * SAMPLE_RATIO)
    indices = np.random.choice(len(X), sample_size, replace=False)
    X = X[indices]
    y = y[indices]
    print(f"采样后的数据形状: {X.shape}")

# 使用float16节省内存
if USE_FLOAT16:
    print("将数据转换为float16以节省内存...")
    X = X.astype(np.float16)

# 对特征进行 L2 归一化
print("正在对特征进行 L2 归一化...")
X = safe_normalize_L2(X)

# 将数据分为训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

print(f"训练集大小: {X_train.shape}")
print(f"测试集大小: {X_test.shape}")

# 重要修改：在使用 KNN 之前准备特征（将 4D 数据转换为 2D）
print("正在准备特征（将 4D 数据转换为 2D 特征矩阵）...")
X_train_flat = prepare_features(X_train)
X_test_flat = prepare_features(X_test)

# 确保转换后的特征也是 NumPy 数组
if isinstance(X_train_flat, list):
    X_train_flat = np.array(X_train_flat)
if isinstance(X_test_flat, list):
    X_test_flat = np.array(X_test_flat)

# 再次使用float16节省内存
if USE_FLOAT16:
    X_train_flat = X_train_flat.astype(np.float16)
    X_test_flat = X_test_flat.astype(np.float16)

print(f"转换后的训练集特征形状: {X_train_flat.shape}")
print(f"转换后的测试集特征形状: {X_test_flat.shape}")

# 使用传统的 KNN 进行比较 - 注意：由于内存限制，我们使用较小的k值范围
print("正在训练传统 KNN 模型...")
# 尝试不同的 k 值，寻找最佳参数
best_k = 1
best_accuracy = 0

# 减少k的范围以节省内存和时间
for k in range(1, 10):
    try:
        # 使用更少的邻居数和算法优化
        knn = KNeighborsClassifier(n_neighbors=k, algorithm='kd_tree', n_jobs=-1)
        # 使用转换后的 2D 特征进行训练
        knn.fit(X_train_flat, y_train)
        
        # 在测试集上进行预测
        y_pred = knn.predict(X_test_flat)
        
        # 计算准确率
        accuracy = accuracy_score(y_test, y_pred)
        print(f"k = {k}, 准确率: {accuracy:.4f}")
        
        # 更新最佳参数
        if accuracy > best_accuracy:
            best_accuracy = accuracy
            best_k = k
    except Exception as e:
        print(f"k = {k} 时出错: {e}")

print(f"最佳 k 值: {best_k}, 最佳准确率: {best_accuracy:.4f}")

# 使用最佳 k 值重新训练模型
best_knn = KNeighborsClassifier(n_neighbors=best_k, algorithm='kd_tree', n_jobs=-1)
best_knn.fit(X_train_flat, y_train)

# 在测试集上进行预测
y_pred = best_knn.predict(X_test_flat)

# 生成详细的分类报告
print("分类报告:")
print(classification_report(y_test, y_pred))

print("训练完成！")